package com.mingserve.common.core.utils;

import org.jtransforms.fft.DoubleFFT_1D;

import java.util.Arrays;

/**
 * 信号处理相关工具类
 * 使用为JTransforms库 (com.github.wendykierp: JTransforms)
 *
 * @author lvwshuai
 * @date 2023-04-04
 */
public class SignalUtils
{
    /**
     * 按采样频率生成采集时间间隔
     *
     * @param n          数据长度
     * @param sampleRate 采样频率
     * @return double[] 时间间隔保留6位小数
     */
    public static double[] intervals(int n, double sampleRate)
    {
        if (sampleRate <= 0) return null;
        double interval = 1.0 / sampleRate; // 每个数据点的间隔
        double[] intervals = new double[n];
        for (int i = 0; i < n; i++) {
            intervals[i] = MathUtils.round(i * interval, 6);
        }
        return intervals;
    }

    /**
     * 生成一个长度为n/2的数组，表示在频率范围[0, fs/2]内每个频率上的幅值
     *
     * @param n  数据长度
     * @param fs 采样频率
     * @return double[] 频率保留2位小数
     */
    public static double[] freq(int n, double fs)
    {
        double[] freq = new double[n / 2+1];
        for (int i = 0; i <= n / 2; i++) {
            freq[i] = MathUtils.round(i * fs / n, 2);
        }
        return freq;
    }

    /**
     * 加速度速整数值/倍数转成浮点数
     *
     * @param vib  加速度整数值
     * @param gain 倍数
     * @return double[] 加速度信号
     */
    public static double[] vibDivideGain(double[] vib, Integer gain)
    {
        if (vib == null) return null;
        if (gain == null || gain == 0) return vib;
        double[] result = new double[vib.length];
        for (int i = 0; i < vib.length; i++) {
            result[i] = vib[i] / gain;
        }
        return result;
    }

    /**
     * 计算坐标轴的加速度值
     *
     * @param adcBits      传感器精度
     * @param fullScaleUse 量程
     * @param vib          振动数据
     * @return double[] 加速度信号 保留5位小数
     */
    public static double[] vibToG(Double adcBits, Double fullScaleUse, double[] vib)
    {
        if (vib == null) return null;
        if (adcBits == null || fullScaleUse == null) return vib;
        double vibAvg = Arrays.stream(vib).average().orElse(0);
        double vibPow = Math.pow(2, adcBits - 1);

        double[] result = new double[vib.length];
        for (int i = 0; i < vib.length; i++) {
            result[i] = MathUtils.round((vib[i] - vibAvg) * fullScaleUse / vibPow, 5);
        }
        return result;
    }

    /**
     * 去直流分量 (DC滤波)
     *
     * @param signal 原始信号
     * @return double[] 原始信号减去平均值
     */
    public static double[] removeDC(double[] signal)
    {
        if (signal == null) return null;
        double mean = Arrays.stream(signal).average().orElse(0);
        return Arrays.stream(signal).map(s -> s - mean).toArray();
    }

    /**
     * 获取振动信号的FFT
     * 使用Apache Commons Math库 (org.apache.commons:commons-math3)FastFourierTransformer进行FFT计算
     *
     * @param vib 振动信号 任意长度
     * @param Fs  采样频率
     * @return double[][]{vib_fft保留5位小数, f保留2位小数}
     */
    public static double[][] fft(double[] vib, double Fs)
    {
        if(vib == null || vib.length == 0) return null;

        vib = removeDC(vib);

        int N = vib.length;

        // 构建复数数组用于 FFT，实部为 vib，虚部为0
        double[] vibComplex = new double[2 * N];
        for (int i = 0; i < N; i++) {
            vibComplex[2 * i] = vib[i];
            vibComplex[2 * i + 1] = 0.0; // 虚部设为0
        }

        // 使用 JTransforms 库进行 FFT
        DoubleFFT_1D fft = new DoubleFFT_1D(N);
        fft.complexForward(vibComplex);

        // 计算频谱幅值和频率
        double[] vib_fft = new double[N / 2 + 1];
        for (int i = 0; i <= N / 2; i++) {
            double real = vibComplex[2 * i];
            double imag = vibComplex[2 * i + 1];
            vib_fft[i] =  MathUtils.round(2.0 * Math.sqrt(real * real + imag * imag) / N,5);
        }

        // 包含奈奎斯特频率点（Fs / 2）
        double[] f = new double[N / 2 + 1];
        for (int i = 0; i <= N / 2; i++) {
            f[i] = MathUtils.round(i * Fs / N,2);
        }

        return new double[][]{vib_fft, f};
    }

    /**
     * 计算振动信息的包络信号
     *
     * @param vib 振动信号
     * @return double[] 包络信号
     */
    public static double[] env(double[] vib)
    {
        if(vib == null || vib.length == 0) return null;

        int N = vib.length;
        double[] vib_h = new double[2 * N];

        // 将原始信号复制到复数数组，偶数索引为实部，奇数索引为虚部
        for (int i = 0; i < N; i++) {
            vib_h[2 * i] = vib[i];
            vib_h[2 * i + 1] = 0.0; // 虚部设为0
        }

        // 使用 JTransforms 库进行 FFT
        DoubleFFT_1D fft = new DoubleFFT_1D(N);
        fft.complexForward(vib_h);

        // 构建Hilbert变换（保留正频率分量）
        for (int i = 1; i < N / 2; i++) {
            vib_h[2 * i] *= 2;
            vib_h[2 * i + 1] *= 2;
        }
        for (int i = N / 2 + 1; i < N; i++) {
            vib_h[2 * i] = 0;
            vib_h[2 * i + 1] = 0;
        }

        // 进行逆FFT
        fft.complexInverse(vib_h, true);

        // 计算包络信号 Math.sqrt(real^2 + imag^2)
        double[] vibEnvelope = new double[N];
        for (int i = 0; i < N; i++) {
            double real = vib_h[2 * i];
            double imag = vib_h[2 * i + 1];
            vibEnvelope[i] = Math.sqrt(real * real + imag * imag);
        }

        return vibEnvelope;
    }
}
